import os
import random

def split_data(base_folder, train_rate=0.9):
    train_file = open(os.path.join(base_folder, 'train.txt'), 'w')
    val_file = open(os.path.join(base_folder, 'val.txt'), 'w')

    # 获取所有类别的子文件夹
    sub_folders = [f for f in os.listdir(base_folder) if os.path.isdir(os.path.join(base_folder, f))]

    for folder in sub_folders:
        folder_path = os.path.join(base_folder, folder)
        images = [img for img in os.listdir(folder_path) if img.lower().endswith(('.png', '.jpg', '.jpeg'))]

        # 打乱图像顺序以随机划分
        random.shuffle(images)
        split_index = int(len(images) * train_rate)

        # 训练数据
        train_images = images[:split_index]
        for img in train_images:
            train_file.write(f"{folder}/{img} {folder}\n")

        # 验证数据
        val_images = images[split_index:]
        for img in val_images:
            val_file.write(f"{folder}/{img} {folder}\n")

    train_file.close()
    val_file.close()

# 使用示例
base_folder = "/Users/tunm/datasets/person_uniform/classify"  # 替换为您的文件夹路径
split_data(base_folder, train_rate=0.9)
